Unified-IO 2：首个图像/文本/音频/视频/动作多模态模型，迈向AGI任务
“Unified-IO 2: Scaling Autoregressive Multimodal Models with Vision, Language, Audio, and Action”
摘要
简介
相关工作
方法
UNIFIED-IO 2是一个多模态任务的模型，使用统一的任务表示方法。模型架构和稳定训练技术是关键，同时使用多模态训练目标和效率优化方法。UNIFIED-IO 2的目标是实现多模态任务的高效处理和准确性能。
统一任务表示
UNIFIED-IO 2通过一个统一的编码器-解码器转换器处理所有模态，通过将各种输入和输出（图像、文本、音频、动作、框等）编码为共享表示空间中的标记序列来实现。我们的编码过程遵循UNIFIED-IO的设计，进行了一些修改以提高性能，并为其他模式提供了新的编码器和解码器。
文本、稀疏结构和动作。文本输入和输出使用来自LLaMA的字节对编码进行标记。稀疏结构（如边界框、关键点和相机姿势）被离散化，然后使用1000个特殊标记添加到词汇表中进行编码。点用两个符号序列（一个代表x，一个代表y）编码，边界框用四个符号序列（左上角和右下角）编码，3D长方体用12个符号表示，这些符号编码投影中心、虚拟深度、对数归一化的盒子维度和连续的非中心旋转。对于嵌入任务，离散的机器人动作被生成为文本命令（例如，“move ahead”命令机器人在导航中向前移动）。使用特殊的标记来编码机器人的状态，例如它的位置和旋转。
图像和密集结构。图像使用预训练的视觉转换器（ViT）进行编码。我们将来自ViT的第二层和倒数第二层的特征连接起来，以捕获低级和高级视觉信息。为了生成图像，我们使用VQ-GAN将图像转换为离散标记。这些标记被添加到词汇表中，然后用作目标输出序列，以生成图像。
我们将每像素标签（包括深度、表面法线和二值分割掩码）表示为RGB图像，这些图像可以用我们的图像生成和编码能力生成或编码。对于分割，UNIFIED-IO 2被训练来预测给定类和边界框的二进制掩码。首先进行检测，然后向模型查询每个检测到的边界框和类的分割掩码，可以对整个图像进行分割。
音频。UNIFIED-IO 2将高达4.08秒的音频编码为频谱图。然后使用预训练的音频频谱图转换器（AST）对频谱图进行编码，通过连接AST的第二层和倒数第二层特征并应用线性层来构建输入嵌入，就像使用图像ViT一样。为了生成音频，我们使用ViT-VQGAN将音频转换为离散标记。
图像和音频历史。我们允许最多四个额外的图像和音频片段作为输入，我们将其称为图像或音频历史。这些元素也使用ViT或AST进行编码，但我们随后使用感知器重采样器，进一步将特征压缩为更少数量的标记（图像32个，音频16个）。这种方法大大减少了序列长度，并允许模型在使用历史中的元素作为上下文的同时，以高水平的细节检查图像或音频段。此历史记录用于对以前的视频帧、以前的音频片段或参考图像进行编码，以用于诸如多视图图像重建或图像条件图像编辑等任务。文本词汇表中添加了8个特殊的标记，用于在文本输入或输出中引用这些历史记录中的各个元素。
架构
UNIFIED-IO 2使用Transformer编码器-解码器架构。然而，我们观察到，随着我们整合其他模式，使用UNIFIED-IO 的标准实施会导致训练越来越不稳定。为了解决这个问题，我们进行了各种架构变化，这些变化显著地稳定了多模式训练。
2D旋转嵌入。我们在每个Transformer层应用旋转位置嵌入(RoPE)代替相对位置嵌入。对于非文本模式，我们将RoPE扩展到二维位置：对于任何二维索引，我们将Transformer注意头的每个查询和键嵌入分成两半，并将由两个坐标中的每一个构建的单独旋转嵌入应用于这两半。
QK归一化。我们观察到，当包括图像和音频模态时，多头注意逻辑值非常大，这导致注意权重变为0或1，并导致训练不稳定。为了解决这个问题，我们在点积注意力计算之前对查询和键应用LayerNorm。
缩放余弦注意力。我们使用感知器重采样器将每个图像帧和音频段压缩为固定数量的标记。我们发现，即使QK归一化，感知器的注意逻辑也会增长到极端值。因此，我们通过使用缩放余弦注意在感知器中应用更严格的归一化，这大大稳定了训练。
联合更新预训练的ViT和AST也可能导致不稳定。因此，我们在预训练期间冻结ViT和AST，并在指令调优结束时对它们进行微调。图4显示，尽管输入和输出模式存在异质性，但我们模型的预训练损失是稳定的。
训练目标
提出了一种通用的多模态预训练方法——多模态混合去噪器（MoD），结合了标准跨度破坏、因果语言建模和极端跨度破坏三种范式。对于文本目标，采用UL2的范式；对于图像和音频目标，定义了两个类似的范式。在训练过程中，使用模态标记和范式标记来指示任务。
本文介绍了一种自回归动态掩蔽的方法，用于图像和音频去噪。传统的自回归方法在解码器端存在信息泄漏问题，但简单地在解码器中掩蔽标记会干扰生成和去噪任务。作者提出的解决方案是，在解码器中掩蔽标记，除了在预测该标记时。此外，作者还使用了行、列和卷积形状的掩蔽稀疏注意力。
高效实现
训练多模态数据时，序列长度变化大，需要使用packing技术将多个例子的tokens打包成一个序列，避免跨例子交叉注意力。由于编码器和解码器不支持packing，因此在transformer编码器-解码器阶段前后进行packing。训练时使用启发式算法将长例子与短例子匹配打包。动态packing可提高训练吞吐量。
优化器
使用Adafactor作为优化器，前5000步进行线性预热，学习率衰减为1/ k。训练时使用β1=0.9和β2=1.0-k-0.8，全局梯度裁剪阈值为1.0。训练3.0M步，其中1.5M用于预训练，1.5M用于指令调整。
多模态数据
UNIFIED-IO 2与之前的工作的一个关键区别是，我们从头开始使用多样化的多模态数据进行模型训练。这需要精心策划高质量的开源多模态数据，用于预训练和指令调整。
预训练数据
预训练数据来自各种来源，涵盖多种模式。
- NLP[33%]。使用了用于训练MPT-7B的公开可用数据集。
- 图片和文字[40%]。文本和图像配对数据来自LAION-400M、CC3M、CC12M和RedCaps，为了帮助训练图像历史模式，我们还使用了OBELICS的交错图像/文本数据。
- 视频和音频[25%]。从各种公共数据集中采集音频和视频数据，包括YT-Temporal-1B、ACAV100M、AudioSet、WebVid-10M、HDVILA-10M和Ego4D。
- 3D &具象化数据[1%]。使用CroCo进行交叉视图生成和去噪。Objaverse用于视图合成。ProcTHOR和Habitat中的随机轨迹，用于下一个动作和框架预测。
- 数据增强[1%]。我们考虑两种类型的数据增强：1）从SAM中自动生成分割数据，训练模型对给定点或边界框的对象进行分割；2） 合成补丁检测数据，任务是让模型列出图像中合成的形状的边界框。我们还训练模型输出图像中补丁的总数，以预训练其计数能力。
预训练样本的构造步骤如下：
1）选择目标模态;
2）选择要保留的其他输入方式;
3）选择目标;
4）根据去噪或生成任务生成随机输入掩码;
5）添加指示任务的前缀标记。
指令微调数据
我们通过结合广泛的监督数据集和任务来构建多模态指令调优数据集，包括220个任务，这些任务来自120多个外部数据集。
- 自然语言[25.0%]。使用来自FlanV2和其他各种指令跟随数据集的混合。
- 图像生成[17.6%]。使用与预训练期间相同的图像和文本对。我们还训练模型通过视图合成、图像编辑、基于分割的图像生成和绘图来生成图像。
- 音频生成[7.5%]。包括文本到音频的数据集，包括野外音频、音乐和人类语音。
- 图像理解[17.8%]。包括各种数据源，包括视觉问答、图像标记、区域分类和具有开放式聊天式响应的数据集。我们还包括多模态指令调谐数据集M3IT和MIMIC-IT。
- 视频理解[10.6%]。包括来自视频字幕、视频标签和视频问答的数据源。我们还使用了来自M3it和MIMIC-IT的例子进行视频教学。
- 音频理解[10.6%]。包括来自音频标注和音频字幕的数据源。我们还包括来自视频动作分类的数据，数据集中有音频。
- 图像稀疏标记[7.25%]。主要考虑目标检测、参考表情、3D检测、相机姿态预测、文本检测和人体关键点。
- 图像密集标记[4.06%]。我们做了几个图像标记任务，包括表面法线估计，深度估计和光流。我们还训练我们的模型进行各种分割任务，包括语义分割、本地化分割和引用表达式分割。
- 视频稀疏标记[3.42%]。我们做视频检测、单目标跟踪和视频动作定位。
- 嵌入式AI[4.33%]。对于VIMA-Bench，我们使用图像输入作为对环境的初始观察，以及提示符中图像或视频的图像历史。我们在模拟和现实环境中添加了具有连续控制的大规模操作数据集。我们还从Habitat Gibson场景中训练PointNav任务。
实验
本文评估了预训练和指令调整模型在需要解析和生成所有模态的广泛任务上的表现，包括图像、视频、音频、文本和动作。没有进行任务特定的微调。
预训练评估
UNIFIED-IO 2在多个任务上展现出了与专门的任务模型相媲美甚至更好的性能。尽管进行了大量的多任务学习，但在HellaSwag任务上的结果表明，UNIFIED-IO 2在语言建模方面的能力介于typical 3B和7B语言模型之间。这可能是因为相比于基于语言的语言模型，该模型看到的标记数量要少得多。
GRIT结果
UNIFIED-IO 2在General Robust Image Task (GRIT) Benchmark上表现出色，超过了之前最好的模型UNIFIED-IO。它在定位、分类、分割和关键点等任务上都取得了进步。在VQA任务上，UNIFIED-IO 2在同源问题上表现更好，但在新源问题上表现稍逊。尽管UNIFIED-IO 2略逊于UNIFIED-IO，但它仍然在指代表达得分上表现出色。UNIFIED-IO 2支持更高质量的图像和文本生成，以及更多的任务和模态，展示了其强大的多任务能力。UNIFIED-IO 2甚至在3亿参数模型上保持了更好的整体性能。消融实验结果显示，UNIFIED-IO 2的所有任务在模型规模增加时都有所改善，表明UNIFIED-IO 2受益于规模的扩大。
生成结果
UNIFIED-IO 2在生成图像、音频和动作输出方面表现出色。在图像生成方面，UNIFIED-IO 2的TIFA得分接近minDALL-E，比其他通用模型如CoDi和Emu高出约10个点。在音频生成方面，UNIFIED-IO 2的能力也很强，得分高于大多数专业模型。在动作生成方面，UNIFIED-IO 2通过VIMA-Bench进行评估，与原始的VIMA策略和VIMAIMG策略相比，平均成功率较高。
视觉语言结果
研究人员评估了视觉语言模型的性能，并将其与其他视觉/语言通用模型进行比较。他们在12个视觉/语言基准测试中展示了UNIFIED-IO2的强大表现，包括在VQA、ScienceQA和TallyQA等任务中超越其他通用模型。在指代表达和字幕生成方面，UNIFIED-IO2也取得了不错的成绩。此外，它在最近提出的评估基准测试中也表现出色。总体而言，UNIFIED-IO2在广泛的能力范围内不会牺牲视觉/语言性能。
视频，音频和其他结果
UNIFIED-IO 2在音频和视频分类、字幕以及视频问答方面表现良好。在Seed-Bench Temporal上，UNIFIED-IO 2的表现超过了BLIP-2和InstructBLIP，超过了8.5个点。在Kinetics-Sounds上，UNIFIED-IO 2的表现也优于仅在该数据集上训练的MBT。在Objectron基准测试上，我们的模型表现良好，与CubeRCNN相似。然而，在nuScenes和Hypersim等多目标3D检测任务中，性能显著下降。在COCO目标检测中，我们的模型在除了’stuff’类别之外的平均精度（AP）达到了47.2，AP50为57.7，AP75为50.0。然而，对于包含许多对象的图像，我们的模型存在困难。在深度估计方面，我们的模型在NYUv2深度数据集上的RMSE为0.623。然而，针对这个任务的微调将RMSE提高到了0.423。我们的模型在对象检测的数据增强方面相对较有限。在其他任务方面，如单目标跟踪、机器人操作的未来状态预测和基于图像的3D视图合成等，我们展示了定性的可视化结果。
限制
- 使用更大版本的图像和音频编码器可以显著提高性能。
- 图像生成的质量不如稳定扩散模型，音频生成限制在约4秒。
- 由于计算资源有限，模型的超参数探索受限。使用更大的批量大小可能会提高性能。
- 模型在深度、视频等模态或需要更专业能力（如3D物体检测）方面可靠性较低，可能是因为这些领域的任务种类有限。
- 提高数据质量可以提高模型性能，但尽管付出了很大努力，人工编写的提示仍然缺乏多样性。在处理新的指令任务时，模型的性能明显下降，而不是在训练时的任务。
总结
研究人员开发了UNIFIED-IO 2，这是第一个能够理解和生成图像、文本、音频和动作的自回归多模态模型。该模型通过广泛的多模态数据进行训练，并通过指令调整进一步优化。研究人员提出了各种架构变化来稳定多模态训练，并提出了多模态混合去噪目标，以有效利用多模态信号。该模型在各种任务中取得了有希望的结果。未来，研究人员计划将UNIFIED-IO 2从编码器-解码器模型扩展到仅解码器模型，并扩大模型规模，提高数据质量和完善整体模型设计。研究得到了Google TPU Research Cloud的支持。